#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import requests
import json
import click
from tqdm import tqdm

HEADERS = {'user-agent': 'ApiPOST Runtime +https://www.apipost.cn', 'content-type': 'application/json'}


@click.group()
def main():
    """
    BNLP导出工具合集

    @author: Yunchao Ling
    @create_time: 2023/03/05
    @update_time: 2023/03/30
    """
    pass


# 根据批次号获取关系标注信息
def get_relation_by_batchId(approved: bool, batchId: str, projectCode: str):
    url = "https://idc.biosino.org/bnlp-api/api/relation/findApprovedByBatchId"
    if not approved:
        url = "https://idc.biosino.org/bnlp-api/api/relation/findByBatchId"
    params = {"batchId": batchId, "projectCode": projectCode}
    r = requests.request("GET", url, headers=HEADERS, params=params)
    if r.status_code == 200:
        return r.text
    else:
        return "error"


# 根据实体号获取实体信息
def get_entity_by_id(entityId: str):
    url = "https://idc.biosino.org/bnlp-api/api/entity/findById"
    params = {"id": entityId}
    r = requests.request("GET", url, headers=HEADERS, params=params)
    if r.status_code == 200:
        return r.text
    else:
        return "error"


# 根据实体标签号获取实体标签信息
def get_entity_by_label(labelId: str):
    url = "https://idc.biosino.org/bnlp-api/api/entity/findLabelById"
    params = {"id": labelId}
    r = requests.request("GET", url, headers=HEADERS, params=params)
    if r.status_code == 200:
        return r.text
    else:
        return "error"


# 根据关系ID号获取关系信息
def get_relation_by_id(relationId: str):
    url = "https://idc.biosino.org/bnlp-api/api/relation/findLabelById"
    params = {"id": relationId}
    r = requests.request("GET", url, headers=HEADERS, params=params)
    if r.status_code == 200:
        return r.text
    else:
        return "error"


# 根据articleId号获取文献信息
def get_article_by_id(articleId: str):
    url = "https://idc.biosino.org/bnlp-api/api/article/findById"
    params = {"articleId": articleId}
    r = requests.request("GET", url, headers=HEADERS, params=params)
    if r.status_code == 200:
        return r.text
    else:
        return "error"


# 根据articleId和段落号获取文本
def get_paragraph_by_id(articleId: str, paragraphId: str):
    source_text = get_article_by_id(articleId)
    if source_text == "error":
        return "error"
    else:
        source_text = json.loads(source_text)
        for item in source_text["data"]["body"]:
            for item2 in item["items"]:
                if item2["id"] == paragraphId:
                    content = item2["content"]
                    return content


def get_structured_entity(entityId: str, sublabel: bool, code: bool):
    source_text = get_entity_by_id(entityId)
    if source_text == "error":
        return "error"
    else:
        source_json = json.loads(source_text)["data"]
        article_text = get_paragraph_by_id(source_json["pmid"], source_json["textId"])
        source_json["text"] = article_text
        if sublabel:
            label_json = json.loads(get_entity_by_label(str(source_json["subLabels"])))
        else:
            label_json = json.loads(get_entity_by_label(str(source_json["labelId"])))
        # print(str(source_json["labelId"]))
        # print(label_json)
        if code:
            source_json["type"] = label_json["data"]["code"]
        else:
            source_json["type"] = label_json["data"]["name"]
        return source_json


def get_full_structured_entity(entityId: str, code: bool):
    source_text = get_entity_by_id(entityId)
    if source_text == "error":
        return "error"
    else:
        source_json = json.loads(source_text)["data"]
        article_text = get_paragraph_by_id(source_json["pmid"], source_json["textId"])
        source_json["text"] = article_text
        label_json = json.loads(get_entity_by_label(str(source_json["labelId"])))
        if code:
            source_json["type"] = label_json["data"]["code"]
        else:
            source_json["type"] = label_json["data"]["name"]
        if "subLabels" in source_json:
            label_json2 = json.loads(get_entity_by_label(str(source_json["subLabels"])))
            if code:
                source_json["subtype"] = label_json2["data"]["code"]
            else:
                source_json["subtype"] = label_json2["data"]["name"]

        # print(str(source_json["labelId"]))
        # print(label_json)

        return source_json


def check_entity_sublabel(entityId: str):
    source_text = get_entity_by_id(entityId)
    if source_text == "error":
        return False
    else:
        source_json = json.loads(source_text)["data"]
        # article_text = get_paragraph_by_id(source_json["pmid"], source_json["textId"])
        # source_json["text"] = article_text
        if "subLabels" not in source_json:
            return False
        else:
            return True


@main.command()
@click.option("--approved", default=True, nargs=1, type=bool, help="仅导出审编后数据，默认True")
@click.option("--sublabel", default=True, nargs=1, type=bool, help="导出sublabel标签，默认True")
@click.option("--code", default=False, nargs=1, type=bool, help="导出code而不是name，默认False")
@click.argument("batch_id", nargs=1)
@click.argument("project_code", nargs=1)
@click.argument("out_file", nargs=1)
def get_structured_relation(approved: bool, sublabel: bool, code: bool, batch_id: str, project_code: str, out_file: str):
    """导出关系审编结果"""
    source_text = get_relation_by_batchId(approved, batch_id, project_code)
    if source_text == "error":
        raise IOError("Web Error.")
    else:
        source_json = json.loads(source_text)
        data = []
        for item in tqdm(source_json["data"]):
            # tqdm.write("subject")
            tqdm.write(item["articleId"])
            if item["subject"][0].startswith("E_"):
                item["subject"] = get_structured_entity(item["subject"][0][2:], sublabel, code)
            else:
                item["subject"] = item["subject"][0]
            # tqdm.write("object")
            # tqdm.write(item["objects"][0])
            if item["objects"][0].startswith("E_"):
                item["objects"] = get_structured_entity(item["objects"][0][2:], sublabel, code)
            else:
                item["objects"] = item["objects"][0]
            # tqdm.write("relation")
            relation_json = json.loads(get_relation_by_id(item["relation"]))
            item["relation"] = relation_json["data"]["name"]
            data.append(item)
            # print(item)
        outfile = open(out_file, "w")
        outfile.write(json.dumps({"data": data}, indent=4) + "\n")
        outfile.flush()
        outfile.close()
        # return json.dumps({"data": data}, indent=4)


@main.command()
@click.option("--approved", default=True, nargs=1, type=bool, help="仅导出审编后数据，默认True")
@click.option("--code", default=False, nargs=1, type=bool, help="导出code而不是name，默认False")
@click.argument("batch_id", nargs=1)
@click.argument("project_code", nargs=1)
@click.argument("out_file", nargs=1)
def get_full_structured_relation(approved: bool, code: bool, batch_id: str, project_code: str, out_file: str):
    """导出完整关系审编结果，包括二级标签"""
    source_text = get_relation_by_batchId(approved, batch_id, project_code)
    if source_text == "error":
        raise IOError("Web Error.")
    else:
        source_json = json.loads(source_text)
        data = []
        for item in tqdm(source_json["data"]):
            # tqdm.write("subject")
            tqdm.write(item["articleId"])
            if item["subject"][0].startswith("E_"):
                item["subject"] = get_full_structured_entity(item["subject"][0][2:], code)
            else:
                item["subject"] = item["subject"][0]
            # tqdm.write("object")
            # tqdm.write(item["objects"][0])
            if item["objects"][0].startswith("E_"):
                item["objects"] = get_full_structured_entity(item["objects"][0][2:], code)
            else:
                item["objects"] = item["objects"][0]
            # tqdm.write("relation")
            relation_json = json.loads(get_relation_by_id(item["relation"]))
            item["relation"] = relation_json["data"]["name"]
            data.append(item)
            # print(item)
        outfile = open(out_file, "w")
        outfile.write(json.dumps({"data": data}, indent=4) + "\n")
        outfile.flush()
        outfile.close()
        # return json.dumps({"data": data}, indent=4)


@main.command()
@click.option("--approved", default=True, nargs=1, type=bool, help="仅导出审编后数据，默认True")
@click.argument("batch_id", nargs=1)
@click.argument("project_code", nargs=1)
def check_sublabel(approved: bool, batch_id: str, project_code: str):
    """检查二级标签"""
    source_text = get_relation_by_batchId(approved, batch_id, project_code)
    if source_text == "error":
        raise IOError("Web Error.")
    else:
        source_json = json.loads(source_text)
        data = set()
        for item in tqdm(source_json["data"]):
            # tqdm.write("subject")
            tqdm.write(item["articleId"])
            if item["subject"][0].startswith("E_"):
                if not check_entity_sublabel(item["subject"][0][2:]):
                    data.add(item["articleId"])
            if item["objects"][0].startswith("E_"):
                if not check_entity_sublabel(item["objects"][0][2:]):
                    data.add(item["articleId"])
        print("共计%d篇存在二级标签不全:" % len(data))
        for item in data:
            print(item)


if __name__ == '__main__':
    main()
    # outfile = open("/Users/genesis/Downloads/ind_230304.json", "w")
    # outfile.write(get_structured_relation("338", "ff6c197e17a848ac83a525bedf22f840") + "\n")
    # outfile.flush()
    # outfile.close()

